import json
import time
from huey.exceptions import TaskException
from langchain_core.documents import Document
from src.build_aitemplate.to_document import to_document_bpr, to_document_sop
from src.build_aitemplate.db import faiss_embedding

from src.build_aitemplate.datamodel import AITemplateModel
from src.build_aitemplate.build_aitemplate import BuildSOPAiTemplate, BuildBPRAiTemplate
from src.huey_init._huey import huey
from src.tools.dms_operator import dmsoperator
from src.tools.excel_to_md import excel2md

@huey.task(context=True)
def run_aitemplate_task(aitemplate_model: AITemplateModel, task=None):
    file_name = aitemplate_model.format_ref_file_list[0].get("documentName")
    # 根据文件名判断是建模工艺规程，还是建模生产批记录
    if "工艺规程" in file_name:
        run_sop_task(aitemplate_model, task)
    elif "记录" in file_name:
        run_bpr_task(aitemplate_model, task)
    else:
        errorinfo = {
            "FunctionType": aitemplate_model.exchange_info.get("functionType"),
            "TaskId": aitemplate_model.exchange_info.get("dmsTaskId"),
            "AiTaskId": task.id,
            "Reason": "没有工艺规程或批记录"
        }
        dmsoperator.post_errortodms(errorinfo=errorinfo)
        raise TaskException("没有工艺规程或批记录")


def run_sop_task(aitemplate_model: AITemplateModel, task):
    aitemp = BuildSOPAiTemplate()

    # 内容引用文件入向量库
    template_files = aitemplate_model.content_cit_file_list
    for template_file in template_files:
        template_file_path = template_file.get("realPath")
        if template_file_path.endswith(('.xls', '.xlsx')):
            suppliertext = excel2md(template_file_path)
            doc = Document(page_content=suppliertext)
            faiss_embedding.add_documents([doc])
        else:
            template_file_jsondata = dmsoperator.get_jsonfile(template_file_path)
            try:
                docus = to_document_sop(template_file_jsondata, template_file)
            except:
                errorinfo = {
                    "FunctionType": aitemplate_model.exchange_info.get("functionType"),
                    "TaskId": aitemplate_model.exchange_info.get("dmsTaskId"),
                    "AiTaskId": task.id,
                    "Reason": "来源文档为空"
                }
                dmsoperator.post_errortodms(errorinfo=errorinfo)
                raise TaskException("来源文档为空")
            faiss_embedding.add_documents(docus)
    retr = faiss_embedding.as_retriever(search_type="mmr", search_kwargs={"k": 3, "fetch_k": 30})

    for format_ref_file in aitemplate_model.format_ref_file_list:
        # 建模工艺规程入口
        template_format_res = aitemp.build_sop_aitemplate(format_ref_file, retr, task.id)
        exchange_info = {
            "functionType": aitemplate_model.exchange_info.get("functionType"),
            "taskType": aitemplate_model.exchange_info.get("taskType"),
            "taskDocumentType": aitemplate_model.exchange_info.get("taskDocumentType"),
            "executionType": aitemplate_model.exchange_info.get("executionType"),
            "isSuccess": True,
            "dmsTaskId": aitemplate_model.exchange_info.get("dmsTaskId"),
            "docId": format_ref_file.get("docId"),
            "taskDocId": format_ref_file.get("taskDocId"),
            "resultFirstFilePath": "string",
            "resultSecondFilePath": "string",
            "aiTaskId": task.id,
            "aiTemplateType": 0,
            "aiTemplate": template_format_res
        }
        # res_json = json.dumps(exchange_info, indent=2, ensure_ascii=False)
        # path = f"{time.time()}ai建模11111.json"
        # with open(path, "w", encoding="utf-8") as f:
        #     f.write(res_json)
        dmsoperator.post_exchange_info(exchange_info)

def run_bpr_task(aitemplate_model: AITemplateModel, task):
    aitemp = BuildBPRAiTemplate()

    # 内容引用文件入向量库
    template_files = aitemplate_model.content_cit_file_list
    for template_file in template_files:
        template_file_path = template_file.get("realPath")
        if template_file_path.endswith(('.xls', '.xlsx')):
            suppliertext = excel2md(template_file_path)
            doc = Document(page_content=suppliertext)
            faiss_embedding.add_documents([doc])
        else:
            template_file_jsondata = dmsoperator.get_jsonfile(template_file_path)
            docus = to_document_bpr(template_file_jsondata, template_file)
            faiss_embedding.add_documents(docus)
    retr = faiss_embedding.as_retriever(search_type="mmr", search_kwargs={"k": 3, "fetch_k": 30})

    for format_ref_file in aitemplate_model.format_ref_file_list:
        # 建模批记录入口
        template_format_res = aitemp.build_bpr_aitemplate(format_ref_file, retr, task.id)
        exchange_info = {
            "functionType": aitemplate_model.exchange_info.get("functionType"),
            "taskType": aitemplate_model.exchange_info.get("taskType"),
            "taskDocumentType": aitemplate_model.exchange_info.get("taskDocumentType"),
            "executionType": aitemplate_model.exchange_info.get("executionType"),
            "isSuccess": True,
            "dmsTaskId": aitemplate_model.exchange_info.get("dmsTaskId"),
            "docId": format_ref_file.get("docId"),
            "taskDocId": format_ref_file.get("taskDocId"),
            "resultFirstFilePath": None,
            "resultSecondFilePath": None,
            "aiTaskId": task.id,
            "aiTemplateType": 1,
            "aiTemplate": template_format_res
        }
        # res_json = json.dumps(exchange_info, indent=2, ensure_ascii=False)
        # path = f"{time.time()}ai建模11111.json"
        # with open(path, "w", encoding="utf-8") as f:
        #     f.write(res_json)
        dmsoperator.post_exchange_info(exchange_info)
